import numpy as np
import casadi as ca

from src.model import Model


class LocalPlanner:
    def __init__(self, model: Model, timestep, N=20):
        super().__init__()

        self._N = N

        self._model = model

        self._opti = ca.Opti()
        self._opt_controls = self._opti.variable(self._N, 2)

        self._opt_states = self._opti.variable(self._N + 1, 3)

        self._opt_x0 = self._opti.parameter(3)
        self._opt_xs = self._opti.parameter(3)

        self._obstacles_num = 5
        self._opt_obstacles = self._opti.parameter(self._obstacles_num, 3)

        for i in range(self._N):
            x_next = self._opt_states[i, :] + self._model.derivatives(self._opt_states[i, :],
                                                                      self._opt_controls[i, :]).T * timestep
            self._opti.subject_to(self._opt_states[i + 1, :] == x_next)

        self._slack_obs = 1e5
        self._slacks = []
        for i in range(1, self._N + 1):
            for j in range(self._obstacles_num):
                temp_constraints = (self._opt_states[i, 0] - self._opt_obstacles[j, 0]) ** 2 + (
                        self._opt_states[i, 1] - self._opt_obstacles[j, 1]) ** 2
                slacks_k = self._opti.variable()
                self._slacks.append(slacks_k)
                self._opti.subject_to(slacks_k >= 0)
                self._opti.subject_to(slacks_k <= 0.05)
                self._opti.subject_to(
                    temp_constraints >= ((self._model.radius + self._opt_obstacles[j, 2]) ** 2) - slacks_k)

        self._Q = np.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 1.0]])
        self._R = np.array([[1.0, 0.0], [0.0, 1.0]])

        self._cost = 0

        self._opti.subject_to(self._opti.bounded(-self._model.v_max, self._opt_controls[:, 0], self._model.v_max))
        self._opti.subject_to(
            self._opti.bounded(-self._model.omega_max, self._opt_controls[:, 1], self._model.omega_max))

        self._opts_setting = {
            'ipopt.max_iter': 100,
            'ipopt.print_level': 0,
            'print_time': 0,
            'ipopt.acceptable_tol': 1e-8,
            'ipopt.acceptable_obj_change_tol': 1e-6
        }
        self._opti.solver('ipopt', self._opts_setting)

        self._u0 = np.zeros((N, 2))
        self._next_states = np.zeros((N + 1, 3))

    def init(self):
        self._opti.subject_to(self._opt_states[0, :] == self._opt_x0.T)

        for i in range(self._N):
            self._cost += ca.mtimes(
                [(self._opt_states[i, :] - self._opt_xs.T), self._Q,
                 (self._opt_states[i, :] - self._opt_xs.T).T]) + ca.mtimes(
                [self._opt_controls[i, :], self._R, self._opt_controls[i, :].T])
        self._cost += self._slack_obs * ca.sumsqr(ca.vertcat(*self._slacks))
        self._opti.minimize(self._cost)

    def compute(self, current_state, desired_state, obstacles):
        self._opti.set_value(self._opt_x0, current_state)
        self._opti.set_value(self._opt_xs, desired_state)
        self._opti.set_value(self._opt_obstacles, obstacles)

        self._opti.set_initial(self._opt_controls, self._u0)
        self._opti.set_initial(self._opt_states, self._next_states)

        sol = self._opti.solve()

        u_res = sol.value(self._opt_controls)
        next_states_pred = sol.value(self._opt_states)

        self._u0[:, :] = np.concatenate((u_res[1:], u_res[-1:]))
        self._next_states[:, :] = np.concatenate((next_states_pred[1:], next_states_pred[-1:]))

        return u_res[0]

    @property
    def next_states(self):
        return self._next_states.copy()
